# from django.test import TestCase
import os

# Create your tests here.
import requests
import json

from twisted.python.util import raises


# 自定义公共函数方法
# AI调用描述：调用代码模板call_third_party_api，去调用第三方接口，请求方法为%s，API的URL为%s，请求头%s，请求参数%s，并返回响应对象的解析数据%s。
def call_third_party_api(method, url, headers=None, params=None, data=None, json_data=None):
    """
    调用第三方 API

    :param method: HTTP 请求方法 (GET, POST, PUT, DELETE 等)
    :param url: API 的 URL
    :param headers: 请求头 (可选)
    :param params: GET请求参数: 查询参数 (可选)
    :param data: POST, PUT, DELETE 等请求参数: 表单数据 (可选)
    :param json_data: POST, PUT, DELETE 等请求参数: JSON 数据 (可选)
    :return: 响应对象
    """
    try:
        response = requests.request(
            method=method,
            url=url,
            headers=headers,
            params=params,
            data=data,
            json=json_data
        )
        # 检查响应状态码
        if response.status_code == 200:
            return response.json()  # 返回 JSON 数据
        else:
            response.raise_for_status()  # 抛出异常，如果响应状态码不是 200
    except requests.RequestException as e:
        print(f"Request failed: {e}")
        return None

def get_workflow_dict():
    workflows_path = r"/workflow/request_prompt.json"
    with open(workflows_path, "r", encoding="utf-8") as f:
        json_data = json.load(f)
    return json_data


def request_comfyui_prompt():
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36 Edg/131.0.0.0"
    }
    method = "POST"
    url = "http://127.0.0.1:8188/api/prompt"
    data = get_workflow_dict()
    json_data = json.dumps(data)
    print(data)
    print("=========================================")
    result = call_third_party_api(method, url, headers=None, params=None, data=json_data, json_data=data)
    print(result)  # 8f9cdfe3-8544-4c09-91ac-4b388ba81120



def get_history(prompt_id=None):
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        }
    method = "GET"
    url = f"http://127.0.0.1:8188/api/history"
    # prompt_id: "bde6e3dc-90b9-46f6-992e-fefabd378cac"
    params = {"max_items": 64}
    if prompt_id:
        url = url+"/{}".format(prompt_id)
        params['prompt_id'] = prompt_id
    result = call_third_party_api(method, url, headers=headers, params={"max_items": 64})
    print(result)
    return result


def get_default_workflow(json_file=None):
    if not json_file or not os.path.exists(json_file):
        # # 获取当前文件的绝对路径
        # current_path = os.path.abspath(__file__)
        # # 获取当前文件所在的目录
        # current_dir = os.path.dirname(current_path)
        base_dir = os.path.abspath(os.path.join(os.getcwd(), ".."))
        print(base_dir)
        # json_file = os.path.join(base_dir, r"\workflow\request_prompt.json")
        # json_file = os.path.join(r'E:\project\python\dashboard\workflow\request_prompt.json')
        json_file = os.path.join(r'G:\ComfyUI\入门工作流\ComfyUI_Workflows\basic\SDXL_simple.json')
        print(json_file)
    with open(json_file, "r", encoding="utf-8") as f:
        json_data = json.load(f)
    return json_data


def workflow_link():
    workflow = get_default_workflow()
    # 二. 替换文本和client_id
    # 1.替换extra_data中的工作流workflow的文本节点内容

    # 先获取到K采样器
    ks_positive_link = list()
    ks_negative_link = list()
    input_link_dict = dict()
    output_link_dict = dict()
    # extra_data = workflow['extra_data']
    # for wf_key, wf_value in extra_data['extra_pnginfo']['workflow'].items():
    for wf_key, wf_value in workflow.items():
        if wf_key == 'nodes':
            for wf_node in wf_value:  # wf_value = [{"id": 6, "type": "CLIPTextEncode", ...}, ]

                for _ in wf_node.get('outputs', []):
                    for _link in _['links']:
                        output_key = '{}-{}'.format(str(wf_node['id']), wf_node['type'])
                        output_link_dict[_link] = output_key

                for _ in wf_node.get('inputs', []):
                    input_key = '{}-{}'.format(str(wf_node['id']), wf_node['type'])
                    input_link_dict[_['link']] = input_key
                    # 获取采样器的正向参数对象和负面参数对象
                    if 'KSampler' in wf_node['type']:
                        if _['name'] == 'positive' and _['type'] == 'CONDITIONING':
                            ks_positive_link.append(_['link'])
                        if _['name'] == 'negative' and _['type'] == 'CONDITIONING':
                            ks_negative_link.append(_['link'])

    if not ks_positive_link:
        raise

    positive_nodes = list()
    negative_nodes = list()
    already_check_nodes = list()

    def get_txt_nodes(_link_num, _text_nodes=[]):
        print(_link_num)
        previous_link_node = output_link_dict.get(_link_num, None)
        _node_type = previous_link_node.split('-')
        already_check_nodes.append(_node_type[0])
        # 若追溯到是K采样器或者追溯过了的节点，则结束后续追溯。
        if 'KSampler' in _node_type[1] and _link_num in already_check_nodes:
            return list(set(_text_nodes))

        if previous_link_node and any(['CLIPTextEncode' in _node_type[1], 'PrimitiveNode' in _node_type[1]]):
            _text_nodes.append(_node_type[0])
            # 继续往前追溯
            keys = [k for k, v in input_link_dict.items() if v == previous_link_node]
            for k in keys:
                _text_nodes.extend(get_txt_nodes(k, _text_nodes))

        return list(set(_text_nodes))

    for _positive_link in ks_positive_link:
        positive_nodes.extend(get_txt_nodes(_positive_link, positive_nodes))

    for _negative_link in ks_negative_link:
        negative_nodes.extend(get_txt_nodes(_negative_link, negative_nodes))
    print('\n')
    print(list(set(positive_nodes)))
    print(list(set(negative_nodes)))


def get_queue():
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36 Edg/131.0.0.0"
    }
    method = "GET"
    url = "http://127.0.0.1:8188/api/queue"
    result = call_third_party_api(method, url, headers=headers)
    print(result)
    return result


if __name__ == '__main__':
    # request_comfyui_prompt()

    # {'prompt_id': 'b1a41f4b-4888-4e30-a990-2c29220356d9', 'number': 3, 'node_errors': {}}
    # prompt_id = "09bb96b4-e1e3-466f-8453-299a80cda68e"
#     # result_data = get_history(prompt_id)
#     # prompt_id_data = result_data.get(prompt_id)
#     # output_data = prompt_id_data.get('outputs')
#     # print(output_data)

    # workflow_link()

    get_queue()
